#include "dicom/storage/encrypt.h"
#include <gwbase/base/string_util.h>
#include <surbasic/AESHelper.h>
#include <surbasic/Digestor.h>
#include <zlib.h>

using namespace SDBasic::aes;

namespace dicom {

Status Encryptor::DoEncryt(const std::string& hexkey, ReaderStream* src,
                           WriterStream* out, HashEntry* hash) {
  Status s;
  dicom::MD5Helper src_md5;
  dicom::SHA1Helper src_sha1;
  dicom::SHA1Helper enc_sha1;
  dicom::SHA512Helper src_sha512;

  size_t src_size = 0;
  size_t enc_size = 0;

  const int kBlockSize = 4096;
  unsigned char buf[kBlockSize];
  Aes256Encrypt cipher(SDBasic::bytearray::ByteArray::fromHexStr(hexkey),
                       kBlockSize);

  int ret, flush;
  unsigned have;
  z_stream strm;
  unsigned char outbuf[kBlockSize];
  strm.zalloc = Z_NULL;
  strm.zfree = Z_NULL;
  strm.opaque = Z_NULL;
  ::deflateInit(&strm, Z_DEFAULT_COMPRESSION);

  while (1) {
    int n = src->Read(buf, sizeof(buf));
    if (n < 0) {
      s = Status(Status::kOther, "read failed");
      goto error;
    }

    if (n > 0) {
      src_md5.Update(buf, n);
      src_sha1.Update(buf, n);
      src_sha512.Update(buf, n);
      src_size += n;
    }

    flush = (n == 0 ? Z_FINISH : Z_NO_FLUSH);
    strm.avail_in = n;
    strm.next_in = buf;

    do {
      strm.avail_out = kBlockSize;
      strm.next_out = outbuf;
      ret = ::deflate(&strm, flush);
      assert(ret != Z_STREAM_ERROR);
      have = kBlockSize - strm.avail_out;
      switch (ret) {
        case Z_NEED_DICT:
        case Z_DATA_ERROR:
        case Z_MEM_ERROR:
          s = Status(Status::kOther, "zip error");
          goto error;
      }

      if (!cipher.Update(outbuf, have)) {
        s = Status(Status::kOther, "aes error");
        goto error;
      }

      if (cipher.len() > 0) {
        int n = out->Write(cipher.data(), cipher.len());
        if (n < 0 || n != cipher.len()) {
          s = Status(Status::kOther, "failed to write");
          goto error;
        }
        enc_sha1.Update(cipher.data(), cipher.len());
        enc_size += cipher.len();
      }
    } while (strm.avail_out == 0);

    if (ret == Z_STREAM_END) break;
  }

  if (!cipher.Finalize()) {
    s = Status(Status::kOther, "AES256.Finalize() failed");
    goto error;
  }

  if (cipher.len() > 0) {
    int n = out->Write(cipher.data(), cipher.len());
    if (n < 0) {
      s = Status(Status::kOther, "failed to write");
      goto error;
    }

    enc_sha1.Update(cipher.data(), cipher.len());
    enc_size += cipher.len();
  }

  if (out->Finalize() < 0) {
    s = Status(Status::kOther, "failed to finalize write");
    goto error;
  }

  src_md5.Finalize();
  src_sha1.Finalize();
  src_sha512.Finalize();
  enc_sha1.Finalize();

  hash->src_md5 = src_md5.digest();
  hash->src_size = src_size;
  hash->src_digest = StringPrintf("%s%s%lx", src_md5.digest().c_str(),
                                  src_sha1.digest().c_str(), src_size);
  hash->src_sha512 = src_sha512.digest();
  hash->hexkey = hexkey;
  hash->enc_size = enc_size;
  hash->enc_digest = enc_sha1.digest();

  ::deflateEnd(&strm);
  return s;

error:
  ::deflateEnd(&strm);
  return s;
}

Status ZipAndEncryptFile(const std::string& hexkey, const std::string& src_file,
                         const std::string& dest_file, HashEntry* hash) {
  Status s;

  FileReader reader;
  s = reader.Open(src_file);
  if (!s.ok()) {
    return s;
  }

  FileWriter writer;
  s = writer.Open(dest_file);
  if (!s.ok()) {
    reader.Close();
    return s;
  }

  Encryptor enc;
  s = enc.DoEncryt(hexkey, &reader, &writer, hash);

  if (s.ok()) {
    assert(hash->IsComplete());
  }

  reader.Close();
  writer.Close();
  return s;
}
}
